// 包含必要的头文件
#include <cuda_runtime.h>
#include <cufft.h>
#include <stdio.h>
#include <stdlib.h>
#include <math.h>
#include <time.h>
#include <string.h>
#include "prof.h"

// 错误检查宏，确保CUDA和cuFFT调用成功
#define CHECK_CUDA_ERROR(call) \
    do { \
        cudaError_t err = (call); \
        if (err != cudaSuccess) { \
            fprintf(stderr, "CUDA错误: %s 在 %s:%d\n", cudaGetErrorString(err), __FILE__, __LINE__); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

#define CHECK_CUFFT_ERROR(call) \
    do { \
        cufftResult_t err = (call); \
        if (err != CUFFT_SUCCESS) { \
            fprintf(stderr, "cuFFT错误: %d 在 %s:%d\n", err, __FILE__, __LINE__); \
            exit(EXIT_FAILURE); \
        } \
    } while (0)

// 常量定义
const int FILTER_ORDER = 256;      // 滤波器阶数
const int SIGNAL_LENGTH = 1024;    // 信号长度，2的幂次方优化FFT性能
const int BATCH_SIZE = 100000;         // 批处理大小，同时处理16组信号

// 生成低通滤波器系数，使用sinc函数模拟理想低通滤波器
void generate_filter_coeffs(float* coeffs, int order, float cutoff) {
    int n = order - 1;
    float wc = cutoff * M_PI; // 截止频率转换为弧度
    for (int i = -n/2; i <= n/2; i++) {
        if (i == 0) {
            coeffs[i + n/2] = wc / M_PI; // 零点处的值
        } else {
            coeffs[i + n/2] = sinf(wc * i) / (M_PI * i); // sinc函数计算系数
        }
    }
}

// 频域相乘的CUDA内核函数，用于批处理多组信号
__global__ void batch_complex_multiply(cufftComplex* signals_freq, cufftComplex* filter_freq, 
                                     cufftComplex* result, int elements_per_signal, int batch_size) {
    int idx = blockIdx.x * blockDim.x + threadIdx.x;
    int signal_idx = idx / elements_per_signal; // 信号索引
    int element_idx = idx % elements_per_signal; // 元素索引
    
    if (signal_idx < batch_size && element_idx < elements_per_signal) {
        // 复数乘法：(a + bi) * (c + di) = (ac - bd) + (ad + bc)i
        float a = signals_freq[signal_idx * elements_per_signal + element_idx].x;
        float b = signals_freq[signal_idx * elements_per_signal + element_idx].y;
        float c = filter_freq[element_idx].x;
        float d = filter_freq[element_idx].y;
        result[signal_idx * elements_per_signal + element_idx].x = a * c - b * d; // 实部
        result[signal_idx * elements_per_signal + element_idx].y = a * d + b * c; // 虚部
    }
}

// 主函数
int main(int argc, char** argv) {
    // 初始化随机数种子
    srand((unsigned int)time(NULL));
    
    // 获取运行次数n
    if (argc != 2) {
        fprintf(stderr, "用法: %s <运行次数>\n", argv[0]);
        return 1;
    }
    int n = atoi(argv[1]);
    if (n <= 0) {
        fprintf(stderr, "运行次数必须为正整数。\n");
        return 1;
    }
    
    profInit();//初始化
    
    profSetSectName(1, "htod 时间");//根据需要设置计时段名称
    profSetSectName(2, "compute_time");
    profSetSectName(3, "dtoh 时间");
    
    // 生成滤波器系数
    float h_filter_coeffs[FILTER_ORDER];
    generate_filter_coeffs(h_filter_coeffs, FILTER_ORDER, 0.1f); // 截止频率设为0.1
    
    // 准备填充后的滤波器系数（补零到SIGNAL_LENGTH）
    float* h_filter_padded = (float*)calloc(SIGNAL_LENGTH, sizeof(float));
    memcpy(h_filter_padded, h_filter_coeffs, FILTER_ORDER * sizeof(float));
    
    // 分配设备内存：滤波器相关
    float* d_filter_padded;
    cufftComplex* d_filter_freq;
    CHECK_CUDA_ERROR(cudaMalloc(&d_filter_padded, SIGNAL_LENGTH * sizeof(float)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_filter_freq, (SIGNAL_LENGTH / 2 + 1) * sizeof(cufftComplex)));
    
    // 将滤波器系数传输到设备
    CHECK_CUDA_ERROR(cudaMemcpy(d_filter_padded, h_filter_padded, SIGNAL_LENGTH * sizeof(float), cudaMemcpyHostToDevice));
    
    // 创建滤波器的FFT计划并执行（只需一次）
    cufftHandle filter_plan;
    CHECK_CUFFT_ERROR(cufftPlan1d(&filter_plan, SIGNAL_LENGTH, CUFFT_R2C, 1));
    CHECK_CUFFT_ERROR(cufftExecR2C(filter_plan, d_filter_padded, d_filter_freq));
    CHECK_CUFFT_ERROR(cufftDestroy(filter_plan));
    
    // 创建信号的批处理FFT和IFFT计划
    cufftHandle batch_fft_plan, batch_ifft_plan;
    CHECK_CUFFT_ERROR(cufftPlan1d(&batch_fft_plan, SIGNAL_LENGTH, CUFFT_R2C, BATCH_SIZE));
    CHECK_CUFFT_ERROR(cufftPlan1d(&batch_ifft_plan, SIGNAL_LENGTH, CUFFT_C2R, BATCH_SIZE));
    
    // 分配信号相关的设备内存
    float* d_signals;
    cufftComplex* d_signals_freq;
    cufftComplex* d_filtered_freq;
    float* d_filtered_signals;
    CHECK_CUDA_ERROR(cudaMalloc(&d_signals, SIGNAL_LENGTH * BATCH_SIZE * sizeof(float)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_signals_freq, (SIGNAL_LENGTH / 2 + 1) * BATCH_SIZE * sizeof(cufftComplex)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_filtered_freq, (SIGNAL_LENGTH / 2 + 1) * BATCH_SIZE * sizeof(cufftComplex)));
    CHECK_CUDA_ERROR(cudaMalloc(&d_filtered_signals, SIGNAL_LENGTH * BATCH_SIZE * sizeof(float)));
    
    // 分配主机内存
    float* h_signals = (float*)malloc(SIGNAL_LENGTH * BATCH_SIZE * sizeof(float));
    float* h_filtered_signals = (float*)malloc(SIGNAL_LENGTH * BATCH_SIZE * sizeof(float));

    // 生成随机信号，范围[-1, 1]
    for (int j = 0; j < SIGNAL_LENGTH * BATCH_SIZE; j++) {
        h_signals[j] = ((float)rand() / RAND_MAX) * 2.0f - 1.0f;
    }
    profEnterSect(0);
    // 主循环：运行n次
    for (int i = 0; i < n; i++) {      
        // 测量htod时间
        profEnterSect(1);
        CHECK_CUDA_ERROR(cudaMemcpy(d_signals, h_signals, SIGNAL_LENGTH * BATCH_SIZE * sizeof(float), cudaMemcpyHostToDevice));
        profLeaveSect(1);

        profEnterSect(2);
        // 执行批处理FFT（实数到复数）
        CHECK_CUFFT_ERROR(cufftExecR2C(batch_fft_plan, d_signals, d_signals_freq));
        
        // 执行频域相乘
        int elements = (SIGNAL_LENGTH / 2 + 1) * BATCH_SIZE;
        int threads_per_block = 256;
        int blocks_per_grid = (elements + threads_per_block - 1) / threads_per_block;
        batch_complex_multiply<<<blocks_per_grid, threads_per_block>>>(d_signals_freq, d_filter_freq, 
                                                                     d_filtered_freq, SIGNAL_LENGTH / 2 + 1, BATCH_SIZE);
        CHECK_CUDA_ERROR(cudaGetLastError());
        
        // 执行批处理IFFT（复数到实数）
        CHECK_CUFFT_ERROR(cufftExecC2R(batch_ifft_plan, d_filtered_freq, d_filtered_signals));
        CHECK_CUDA_ERROR(cudaDeviceSynchronize());
        profLeaveSect(2);
        // 测量dtoh时间
        profEnterSect(3);
        CHECK_CUDA_ERROR(cudaMemcpy(h_filtered_signals, d_filtered_signals, 
                                  SIGNAL_LENGTH * BATCH_SIZE * sizeof(float), cudaMemcpyDeviceToHost));
        profLeaveSect(3);
    }
    profLeaveSect(0);

    profPrintInfo();
    
    // 清理资源
    free(h_filter_padded);
    free(h_signals);
    free(h_filtered_signals);
    
    CHECK_CUDA_ERROR(cudaFree(d_filter_padded));
    CHECK_CUDA_ERROR(cudaFree(d_filter_freq));
    CHECK_CUDA_ERROR(cudaFree(d_signals));
    CHECK_CUDA_ERROR(cudaFree(d_signals_freq));
    CHECK_CUDA_ERROR(cudaFree(d_filtered_freq));
    CHECK_CUDA_ERROR(cudaFree(d_filtered_signals));
    CHECK_CUFFT_ERROR(cufftDestroy(batch_fft_plan));
    CHECK_CUFFT_ERROR(cufftDestroy(batch_ifft_plan));
    
    return 0;
}
